#!/usr/bin/env python
from astropy.modeling.fitting import LevMarLSQFitter
from astropy.modeling.functional_models import Const2D
from astropy.modeling.models import Gaussian2D
from astropy.table import Table
from scipy.ndimage import binary_erosion, binary_dilation
from scipy.stats import sigmaclip, trim_mean
from skimage.morphology import label
from astropy.io import fits
import sys
import numpy as np
from glob import glob
from argparse import ArgumentParser
import os

def get_bkg(image):
    clip = sigmaclip(image, 7, 7)[0]
    bkg = clip.mean()
    sigma = clip.std()
    index = image < bkg + sigma * 3
    index = binary_erosion(index, iterations=2)
    clip = sigmaclip(image[index], 5, 5)[0]
    bkg = clip.mean()
    sigma = clip.std()
    return bkg, sigma

def get_segmentation(image, det_thresh=10.0, clean_erosion=2, clean_dilation=0,
                     peak_thresh=None, value_thresh=None, npix_thresh=None,
                     low_limit=-9999, largest=False, mask=None):
    # background
    bkg, sigma = get_bkg(image)
    img = image - bkg

    # detect source
    index = img > sigma * det_thresh
    if mask is not None:
        index = np.logical_and(index, mask == 0)
    if clean_erosion > 0:
        index = binary_erosion(index, iterations=clean_erosion)
    if clean_erosion + clean_dilation > 0:
        index = binary_dilation(index, iterations=clean_erosion + clean_dilation)
    if clean_dilation > 0:
        index = binary_erosion(index, iterations=clean_dilation)
    segmap, n = label(index, return_num=True)
    if n == 0:
        return segmap

    # get threshold
    thresh = None
    if peak_thresh is not None:
        if mask is None:
            vmax = np.max(img)
        else:
            vmax = np.max(img[mask == 0])
        thresh = peak_thresh * vmax
    if value_thresh is not None:
        if thresh is None:
            thresh = value_thresh
        else:
            thresh = max(thresh, value_thresh)

    # filter objects
    if thresh is not None or npix_thresh is not None:
        for i in range(1, n + 1):
            index = segmap == i
            if thresh is not None and np.max(img[index]) < thresh:
                segmap[index] = 0
            if npix_thresh is not None and np.sum(index) < npix_thresh:
                segmap[index] = 0
        segmap, n = label(segmap > 0, return_num=True)
    if n == 0:
        return segmap

    # find the largest object
    if largest is True and n > 1:
        segs, n = np.unique(segmap[segmap > 0], return_counts=True)
        seg = segs[np.argmax(n)]
        segmap = np.asarray(segmap == seg).astype(int)

    # remove all low flux pixels
    index = img < low_limit
    segmap[index] = 0

    return segmap


def measure_coord(image, segmap):
    segs = np.unique(segmap[segmap > 0])
    ny, nx = image.shape
    yy, xx = np.mgrid[0:ny, 0:nx]
    x, y, xmin, xmax, ymin, ymax = [], [], [], [], [], []
    for seg in segs:
        index = segmap == seg
        data = image[index]
        x.append((data * xx[index]).sum() / data.sum())
        y.append((data * yy[index]).sum() / data.sum())
        xmin.append(np.min(xx[index]).astype(int))
        xmax.append(np.max(xx[index]).astype(int))
        ymin.append(np.min(yy[index]).astype(int))
        ymax.append(np.max(yy[index]).astype(int))
    tab = Table()
    tab['objid'] = segs
    tab['x'] = x
    tab['y'] = y
    tab['xmin'] = xmin
    tab['xmax'] = xmax
    tab['ymin'] = ymin
    tab['ymax'] = ymax
    return tab


def measure_flux(image, segmap, bkgthick=10, bkggap=10):
    segs = np.unique(segmap[segmap > 0])
    pixels, fluxes, errors, bkgs, noises = [], [], [], [], []
    for seg in segs:
        index = segmap == seg
        npix = np.sum(index)
        index_avoid = binary_dilation(index, iterations=bkggap)
        index_bkg = binary_dilation(index_avoid, iterations=bkgthick)
        index_bkg = np.logical_and(index_bkg, segmap == 0)
        bkg = trim_mean(image[index_bkg], 0.2)
        noise = sigmaclip(image[index_bkg], 7, 7)[0].std()
        flux = image[index].sum() - bkg * npix
        error = noise * np.sqrt(npix)
        pixels.append(npix)
        fluxes.append(flux)
        errors.append(error)
        bkgs.append(bkg)
        noises.append(noise)
    tab = Table()
    tab['objid'] = segs
    tab['npix'] = pixels
    tab['flux'] = fluxes
    tab['fluxerr'] = errors
    tab['bkg'] = bkgs
    tab['noise'] = noises
    return tab

def fit_gauss(image, maxiter=100, fix_const=False):
    ny, nx = image.shape
    yy, xx = np.mgrid[:ny, :nx]
    index = image - np.max(image) > -1e-6
    x0 = xx[index].mean()
    y0 = yy[index].mean()
    amp = image[index].mean()
    fitter = LevMarLSQFitter()
    model = Gaussian2D(x_mean=x0, y_mean=y0, amplitude=amp) + Const2D(amplitude=0)
    model.fixed['amplitude_1'] = fix_const
    index = binary_dilation(index, iterations=20)
    fit = fitter(model, xx[index], yy[index], image[index], maxiter=maxiter)
    return fit, fitter.fit_info

def get_spot_sigma(image):
    # Gaussian fit to get sigma/fwhm
    fit_res, fit_info = fit_gauss(image, fix_const=True)
    sigma = np.sqrt((fit_res.x_stddev_0.value ** 2 + fit_res.y_stddev_0.value ** 2) / 2)
    return sigma


def get_ellip(x, y, flux=1.0, weight=1.0, x0=None, y0=None):
    f = np.sum(flux * weight)
    if x0 is None or y0 is None:
        x0 = np.sum(x * flux * weight) / f
        y0 = np.sum(y * flux * weight) / f
    qxx = np.sum(flux * weight * (x - x0) ** 2) / f
    qyy = np.sum(flux * weight * (y - y0) ** 2) / f
    qxy = np.sum(flux * weight * (x - x0) * (y - y0)) / f
    q1 = (qxx - qyy) / (qxx + qyy)
    q2 = qxy * 2 / (qxx + qyy)
    ellp = np.sqrt(q1 ** 2 + q2 ** 2)
    return ellp

def get_spot_photometry(image, sigma=None, radius=None, radius_factor=4):
    if sigma is None:
        sigma = get_spot_sigma(image)
    if radius is None:
        radius = sigma * 2.355 * radius_factor
        if radius < 3:
            radius = 3

    # iteratively determine the centroid
    ny, nx = image.shape
    yy, xx = np.mgrid[:ny, :nx]
    index = np.max(image) - image < 1e-6
    index = binary_dilation(index, iterations=int(radius) + 1)
    x0, y0 = -1, -1
    for i in range(10):
        x1 = np.sum(xx[index] * image[index]) / np.sum(image[index])
        y1 = np.sum(yy[index] * image[index]) / np.sum(image[index])
        if x1 == x0 and y1 == y0:
            break
        x0, y0 = x1, y1
        index = (xx - x0) ** 2 + (yy - y0) ** 2 < radius ** 2

    # calculate ellipticity
    gauss = Gaussian2D(amplitude=1.0, x_mean=x0, y_mean=y0,
                       x_stddev=sigma, y_stddev=sigma)
    weight = gauss(xx, yy)
    ellp = get_ellip(xx[index], yy[index], flux=image[index],
                    weight=weight[index], x0=x0, y0=y0)
 
    flux = image[index].sum()
    return x0, y0, sigma, ellp, flux

def _write_log(flog, fname, result, bias_flag):
    result_str = []
    for i in range(len(result)):
        result_str.append(str(result[i]))
    if os.path.exists(flog) is False:
        with open(flog, 'w') as f:
            f.write('# name, bias_sub, x, y, R80, ellipticity, flux\n')
    with open(flog, 'a') as f:
        f.write(fname + ', ' + str(bias_flag) +', '+ ', '.join(result_str) + '\n')

def _run(dir_list, flog, dirflag=True, fbias=None, sort=None):
   #print(fbias is not None)
   #print(dir_list)
   if fbias is not None:
       bias_flag = True
       if not os.path.exists(fbias):
           print("bias "+fbias+" does not exist")
           return
       bias = fits.getdata(fbias)
   else:
       bias_flag = False
   if dirflag:
       files = glob(dir_list+"/scan*.fits") 
   else:
       files = dir_list
   if sort is not None:
       idx = []
       for i in range(len(files)):
           idx.append(int(files[i][:-5].split("-")[sort]))
       files = np.array(files)[np.argsort(np.array(idx))].tolist()
   for i in range(len(files)):
       f = files[i]
       if not os.path.exists(f):
           print("file "+f+" does not exist, skip")
           continue
       data = fits.getdata(f).astype("float")    
       if fbias is not None:
           data -= bias
       data -= get_bkg(data)[0]
       fit_res, fit_info = fit_gauss(data, fix_const=True)
       gauss_sigma = np.sqrt((fit_res.x_stddev_0.value ** 2 + fit_res.y_stddev_0.value ** 2) / 2)
       res = get_spot_photometry(data, sigma=gauss_sigma)
       fwhm = res[2]*2.355
       R80 = fwhm
       print("file",f.split('/')[-1],"x =",res[0],"y =",res[1],"R80 = ", res[2]*2.355*0.762, "ellipse =",res[3], "flux =",res[4])
       if flog is not None:
           #print(f.split('/')[-1])
           res1 = (res[0],res[1],res[2]*2.355*0.762,res[3],res[4])
           _write_log(flog,f.split('/')[-1],res1,bias_flag)
       
if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument('--flist', default=None, type=str, nargs='*', help='input files')
    parser.add_argument('--dir', default=None, type=str, help='input directory')
    parser.add_argument('--log', default=None, type=str,
                        help='output result table file')
    parser.add_argument('--bias', default=None, type=str,
                        help='bias file name')
    parser.add_argument('--sort', default=None, type=int,
                        help='sort files with names in format scan-*-*.fits according to which index')
    pars = parser.parse_args()
    if pars.dir is not None:
        _run(pars.dir, pars.log, True, pars.bias, pars.sort)
    else:
        _run(pars.flist, pars.log, False, pars.bias, pars.sort)
